from abc import ABC, abstractmethod
from .runnable import Runnable
from typing import Type, Generic, TypeVar, Dict, Any, Callable, Optional
from pydantic import BaseModel
import logging 
logger = logging.getLogger(__name__)

class TaskInfo(BaseModel):
    runnable: Runnable
    task_name: str

    class Config:
        arbitrary_types_allowed = True

    @classmethod
    def create_task_info(cls, runnable: Runnable, task_name: str):
        return cls(runnable = runnable, task_name = task_name)

class Task(ABC):
    def __init__(self):
        self._ctx: Dict[str, Any] = {}
        self.task_info = None

    @abstractmethod
    def execute(self):
        pass

    @abstractmethod
    def dispose(self):
        """Dispose of any resources used by the task. Considered a destructor."""
        pass

    def get_ctx(self):
        return self._ctx

    def set_ctx(self, ctx):
        self._ctx = ctx

    def set_context(self, key: str, value: Any):
        """Set a value in the task's context."""
        self._ctx[key] = value

    def get_context(self, key: str, default=None) -> Any:
        """Get a value from the task's context."""
        return self._ctx.get(key, default)

T = TypeVar('T', bound='Task')

class TaskFactory(ABC, Generic[T]):
    @abstractmethod
    def create_task(self, info: TaskInfo) -> T:
        """Create a task instance with the given runnable."""
        pass



class DefaultRunnableTask(Task):
    RUNNABLE_ID_IN_CTX_KEY = "current_runnable_id"
    def __init__(self, taskinfo: TaskInfo, pre_run: Optional[Callable[['DefaultRunnableTask'], None]] = None, post_run: Optional[Callable[['DefaultRunnableTask', Optional[Any]], None]] = None):
        super().__init__()
        self.task_info = taskinfo
        self.runnable = taskinfo.runnable
        self.pre_run = pre_run
        self.post_run = post_run

    def add_pre_run_callable(self, pre_run: Callable[['DefaultRunnableTask'], None]):
        self.pre_run = pre_run

    def add_post_run_callable(self, post_run: Callable[['DefaultRunnableTask', Optional[Any]], None]):
        self.post_run = post_run

    def execute(self):
        # if there's ctx in the task, set it to the runnable
        logger.info("Task %s with runnable %s is start exectuiong", self.task_info.task_name, self.runnable.runnable_name)
        ret = None
        try:
            self.set_context(DefaultRunnableTask.RUNNABLE_ID_IN_CTX_KEY, self.runnable.runnable_name)
            
            if self.pre_run and callable(self.pre_run):
                logger.debug("Task %s with runnable %s is executing pre-run logic", self.task_info.task_name, self.runnable.runnable_name)
                self.pre_run(self)  # Execute pre-run logic with task instance
            else:
                logger.warning("Pre-run is not callable for task %s with runnable %s", self.task_info.task_name, self.runnable.runnable_name)

            ret = self.runnable.execute()
            if isinstance(ret, BaseModel):
                ret = ret.model_dump()
            elif isinstance(ret, dict):
                pass
            else:
                raise ValueError(f"Result of runnable {self.runnable.runnable_name} is not a pydantic BaseModel or not a dict, type is {type(ret)}")

            if self.post_run and callable(self.post_run):
                logger.debug("Task %s with runnable %s is executing post-run logic", self.task_info.task_name, self.runnable.runnable_name)
                self.post_run(self, ret)  # Execute post-run logic with task instance and execution result
            else:
                logger.warning("Post-run is not callable for task %s with runnable %s", self.task_info.task_name, self.runnable.runnable_name)
            self.set_context(DefaultRunnableTask.RUNNABLE_ID_IN_CTX_KEY, "")
            if isinstance(ret, dict):
                logger.debug(f"task {self.task_info.task_name} execute finished, result keys: {ret.keys()}")
            # not return ret to avoid cloudpickle error ,just use post_run to handle the result for example: use ctx to store the result
            return ret

        except Exception as err:
            self.set_context(DefaultRunnableTask.RUNNABLE_ID_IN_CTX_KEY, "")
            err = Exception("Failed to run task %s with runnable %s: %s" % (self.task_info.task_name, self.runnable.runnable_name, err))
            raise err

    def dispose(self):
        pass


